// com/tidu/disruptor/manager/BatchAckManager.java
package com.tidu.disruptor.manager;

import com.tidu.disruptor.event.DataSyncEvent;
import com.tidu.strategy.abstractSync.AbstractSyncStrategy;
import com.tidu.strategy.sync.SyncStrategy;
import lombok.extern.slf4j.Slf4j;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentSkipListMap;
import java.util.concurrent.atomic.AtomicLong;

@Slf4j
public class BatchAckManager {
    // 待处理的批次
    private final ConcurrentSkipListMap<Long, DataSyncEvent> pendingBatches = new ConcurrentSkipListMap<>();
    // 失败的批次（用于内存中快速访问）
    private final ConcurrentHashMap<Long, DataSyncEvent> failedBatches = new ConcurrentHashMap<>();
    // 下一个要处理的offset（总是向前推进）
    private final AtomicLong nextOffset = new AtomicLong(0);
    // 同步策略
    private final SyncStrategy syncStrategy;
    // 数据库连接用于记录失败批次
    private final Connection sourceConn;
    private final Connection targetConn;

    public BatchAckManager(SyncStrategy syncStrategy) {
        this.syncStrategy = syncStrategy;
        try {
            this.sourceConn = ((AbstractSyncStrategy)syncStrategy).getSourceDataSource().getConnection();
            this.targetConn = ((AbstractSyncStrategy)syncStrategy).getTargetDataSource().getConnection();
        } catch (Exception e) {
            log.error("Failed to get database connection", e);
            throw new RuntimeException("Failed to get database connection", e);
        }
        // 从存储中加载上次确认的offset
        try {
            Long savedOffset = syncStrategy.loadOffset();
            if (savedOffset != null && savedOffset > 0) {
                nextOffset.set(savedOffset);
                log.info("Loaded last offset: {}", savedOffset);
            }
        } catch (Exception e) {
            log.warn("Failed to load last confirmed offset", e);
        }
    }
    /**
     * 注册新批次
     */
    public void registerBatch(DataSyncEvent event) {
        pendingBatches.put(event.getBatchId(), event);
        log.info("Registered batch {} with offset {} and size {}",
                event.getBatchId(), event.getOffset(), event.getBatchSize());
    }
    /**
     * 简化版本的确认方法（向后兼容）
     */
    public void acknowledgeBatch(long batchId, boolean success) {
        acknowledgeBatch(batchId, success, "Unknown error");
    }
    /**
     * 确认批次处理完成
     */
    public void acknowledgeBatch(long batchId, boolean success, String failReason) {
        DataSyncEvent event = pendingBatches.get(batchId);
        if (event != null) {
            event.markProcessed(success);

            if (success) {
                log.debug("Batch {} processed successfully, offset: {}, size: {}",
                        batchId, event.getOffset(), event.getBatchSize());
            } else {
                log.info("Batch {} processing failed, reason: {}", batchId, failReason);
                log.warn("Batch {} processing failed, offset: {}, size: {}",
                        batchId, event.getOffset(), event.getBatchSize());
                failedBatches.put(batchId, event);
                // 记录失败批次到数据库
                recordFailedBatch(event, failReason);
            }

            // 无论成功失败都更新offset（总是向前推进）
//            updateOffset(event);
            // 尝试更新offset
            updateOffsetIfPossible();
        }
    }
    /**
     * 更新offset（总是向前推进）
     */
    private void updateOffset(DataSyncEvent processedEvent) {
        long currentNextOffset = nextOffset.get();
        long expectedOffset = processedEvent.getOffset();

        // 确保我们处理的是下一个预期的批次
        if (expectedOffset == currentNextOffset) {
            long newNextOffset = expectedOffset + processedEvent.getBatchSize();

            if (nextOffset.compareAndSet(currentNextOffset, newNextOffset)) {
                syncStrategy.saveOffset(newNextOffset);
                log.info("Updated next offset to {} after processing batch {} ({} records)",
                        newNextOffset, processedEvent.getBatchId(), processedEvent.getBatchSize());
                // 从待处理列表中移除
                pendingBatches.remove(processedEvent.getBatchId());
            }
        } else {
            log.warn("Batch {} has offset {} but expected {}, will still advance offset after all pending batches",
                    processedEvent.getBatchId(), expectedOffset, currentNextOffset);

            // 即使不是预期批次，也要确保最终offset得到更新
            pendingBatches.remove(processedEvent.getBatchId());
            // 检查是否所有之前的批次都已处理，如果是则更新offset
            updateOffsetIfNeeded();
        }
    }

    /**
     * 检查并更新offset（当有批次完成时调用）
     */
    private void updateOffsetIfNeeded() {
        long currentNextOffset = nextOffset.get();
        boolean updated = false;

        // 按offset排序检查待处理批次
        for (Map.Entry<Long, DataSyncEvent> entry : pendingBatches.entrySet()) {
            DataSyncEvent event = entry.getValue();

            // 如果批次已经处理完成
            if (event.isProcessed()) {
                // 如果是下一个预期的批次
                if (event.getOffset() == currentNextOffset) {
                    long newNextOffset = event.getOffset() + event.getBatchSize();
                    currentNextOffset = newNextOffset;
                    updated = true;

                    // 移除已处理的批次
                    pendingBatches.remove(entry.getKey());
                    if (!event.isSuccess()) {
                        failedBatches.remove(entry.getKey());
                    }
                } else if (event.getOffset() < currentNextOffset) {
                    // 已经过期的批次，直接移除
                    pendingBatches.remove(entry.getKey());
                    if (!event.isSuccess()) {
                        failedBatches.remove(entry.getKey());
                    }
                }
            }
        }

        // 如果有更新，保存新的offset
        if (updated) {
            if (nextOffset.compareAndSet(nextOffset.get(), currentNextOffset)) {
                syncStrategy.saveOffset(currentNextOffset);
                log.info("Updated next offset to {} after processing pending batches", currentNextOffset);
            }
        }
    }
    /**
     * 检查并更新offset
     *
     * 该方法用于按顺序检查待处理批次的状态，并在满足条件时更新下一个待处理的offset。
     * 只有当批次处理成功且offset是连续的，才会更新offset；
     * 处理失败的批次会被跳过，但会阻止后续批次更新offset（保证顺序性）。
     */
    private void updateOffsetIfPossible() {
        long currentNextOffset = nextOffset.get();
        boolean updated = false;

        // 按顺序检查批次是否已完成
        for (Map.Entry<Long, DataSyncEvent> entry : pendingBatches.entrySet()) {
            DataSyncEvent event = entry.getValue();

            // 如果批次未处理完成，停止检查
            if (!event.isProcessed()) {
                log.debug("Batch {} is not processed yet, stopping check", entry.getKey());
                break;
            }

            // 检查批次是否从期望的位置开始
            if (event.getOffset() == currentNextOffset) {
                /*if (event.isSuccess()) {
                    // 如果批次处理成功，更新下一个offset
                    long newNextOffset = event.getOffset() + event.getBatchSize();
                    currentNextOffset = newNextOffset;
                    updated = true;

                    // 从待处理列表中移除
                    pendingBatches.remove(entry.getKey());
                    log.info("Confirmed batch {} with offset range {}-{}, new next offset: {}",
                            entry.getKey(), event.getOffset(), event.getOffset() + event.getBatchSize(), newNextOffset);
                } else {
                    // 处理失败的批次，不能继续更新offset
                    log.warn("Batch {} at offset {} failed, cannot advance offset",
                            entry.getKey(), event.getOffset());
                    break;
                }*/
                //无论是否成功，都要更新offset，从待处理列表中移除
                long newNextOffset = event.getOffset() + event.getBatchSize();
                currentNextOffset = newNextOffset;
                updated = true;

                // 从待处理列表中移除
                pendingBatches.remove(entry.getKey());
                log.info("Confirmed batch {} with offset range {}-{}, new next offset: {}",
                        entry.getKey(), event.getOffset(), event.getOffset() + event.getBatchSize(), newNextOffset);
            } else if (event.getOffset() > currentNextOffset) {
                // 存在间隙，不能继续处理
                log.debug("Gap detected at batch {}, offset {} > expected {}",
                        entry.getKey(), event.getOffset(), currentNextOffset);
                break;
            } else {
                // offset小于当前nextOffset，说明已经被处理过了，移除
                log.debug("Batch {} at offset {} already processed, removing",
                        entry.getKey(), event.getOffset());
                pendingBatches.remove(entry.getKey());
            }
            // 如果有更新，保存新的offset
            if (updated) {
                if (nextOffset.compareAndSet(nextOffset.get(), currentNextOffset)) {
                    syncStrategy.saveOffset(currentNextOffset); // 保存下一个起始offset
                    log.info("Updated next offset to {}", currentNextOffset);
                }
            }
        }
    }

    /**
     * 记录失败批次到数据库
     */
    private void recordFailedBatch(DataSyncEvent event, String failReason) {
        String sql = "INSERT INTO sync_failed_batches (table_name, batch_id, offset_value, batch_size, fail_reason) VALUES (?, ?, ?, ?, ?)";
        try (PreparedStatement ps = sourceConn.prepareStatement(sql)) {
            // 假设可以从syncStrategy获取表名，或者通过其他方式传递
            String tableName = getTableNameFromSyncStrategy();
            ps.setString(1, tableName);
            ps.setLong(2, event.getBatchId());
            ps.setLong(3, event.getOffset());
            ps.setInt(4, event.getBatchSize());
            ps.setString(5, failReason);

            ps.executeUpdate();
            log.debug("Recorded failed batch {} to database", event.getBatchId());
        } catch (SQLException e) {
            log.error("Failed to record failed batch {} to database", event.getBatchId(), e);
        }
    }
    /**
     * 从SyncStrategy获取表名的辅助方法
     */
    private String getTableNameFromSyncStrategy() {
        // 根据实际的SyncStrategy实现来获取表名
        // 这里可能需要在SyncStrategy接口中添加getTableName()方法
        try {
            if (syncStrategy instanceof com.tidu.strategy.SingleTableSyncDisruptorStrategy) {
                return ((com.tidu.strategy.SingleTableSyncDisruptorStrategy) syncStrategy).getTableName();
            }
        } catch (Exception e) {
            log.warn("Failed to get table name from sync strategy", e);
        }
        return "unknown_table";
    }
    /**
     * 重试失败的批次
     */
    public void retryFailedBatches() {
        String sql = "SELECT batch_id, offset_value, batch_size FROM sync_failed_batches WHERE status = 'PENDING' AND table_name = ?";
        try (PreparedStatement ps = targetConn.prepareStatement(sql)) {
            ps.setString(1, getTableNameFromSyncStrategy());
            ResultSet rs = ps.executeQuery();

            int retryCount = 0;
            while (rs.next()) {
                long batchId = rs.getLong("batch_id");
                long offset = rs.getLong("offset_value");
                int batchSize = rs.getInt("batch_size");

                retryCount++;
                log.info("Scheduled retry for failed batch {} at offset {}", batchId, offset);
            }

            log.info("Scheduled retry for {} failed batches", retryCount);
        } catch (SQLException e) {
            log.error("Failed to query failed batches for retry", e);
        }
    }
    /**
     * 获取下一个待处理的offset（用于获取下一批数据）
     */
    public long getNextOffset() {
        return nextOffset.get();
    }
    /**
     * 获取失败的批次数量
     */
    public int getFailedBatchCount() {
        return failedBatches.size();
    }
    /**
     * 获取待处理的批次数量
     */
    public int getPendingBatchCount() {
        return pendingBatches.size();
    }
}
